import torch
import torch.nn as nn
import torch.nn.functional as nnf
from transformers import GPT2Tokenizer, GPT2LMHeadModel, AdamW, get_linear_schedule_with_warmup
from typing import Tuple, Optional, Union
from enum import Enum

class MappingType(Enum):
    MLP = 'mlp'
    Transformer = 'transformer'


class MLP(nn.Module):

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.model(x)

    def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh):
        super(MLP, self).__init__()
        layers = []
        for i in range(len(sizes) - 1):
            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias))
            if i < len(sizes) - 2:
                layers.append(act())
        self.model = nn.Sequential(*layers)


class MlpTransformer(nn.Module):
    def __init__(self, in_dim, h_dim, out_d: Optional[int] = None, act=nnf.relu, dropout=0.):
        super().__init__()
        out_d = out_d if out_d is not None else in_dim
        self.fc1 = nn.Linear(in_dim, h_dim)
        self.act = act
        self.fc2 = nn.Linear(h_dim, out_d)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.dropout(x)
        return x

class MultiHeadAttention(nn.Module):

    def __init__(self, dim_self, dim_ref, num_heads, bias=True, dropout=0.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim_self // num_heads
        self.scale = head_dim ** -0.5
        self.to_queries = nn.Linear(dim_self, dim_self, bias=bias)
        self.to_keys_values = nn.Linear(dim_ref, dim_self * 2, bias=bias)
        self.project = nn.Linear(dim_self, dim_self)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, y=None, mask=None):
        """多头交叉注意力， x作为查询， y作为键、值， 通过y来指导x关注哪些部分"""
        y = y if y is not None else x
        b, n, c = x.shape
        _, m, d = y.shape
        # b n h dh
        queries = self.to_queries(x).reshape(b, self.num_heads, n, c // self.num_heads)
        # b m 2 h dh
        keys_values = self.to_keys_values(y).reshape(b, m, 2, self.num_heads, c // self.num_heads)
        keys, values = keys_values[:, :, 0].transpose(2, 1), keys_values[:, :, 1].transpose(2, 1)
        # 获得注意力权重，最后两个维度时n, m，把h移到前面，最后n、m就是表明x和y两个序列的位置关系
        attention = torch.einsum('bhnd,bhmd->bhnm', queries, keys) * self.scale
        if mask is not None:
            if mask.dim() == 2:
                mask = mask.unsqueeze(1)
            attention = attention.masked_fill(mask.unsqueeze(3), float("-inf"))
        attention = attention.softmax(dim=2)
        out = torch.einsum('bhnm,bhmd->bhnd', attention, values).reshape(b, n, c)
        out = self.project(out)
        return out, attention


class TransformerLayer(nn.Module):

    def forward_with_attention(self, x, y=None, mask=None):
        x_, attention = self.attn(self.norm1(x), y, mask)
        x = x + x_
        x = x + self.mlp(self.norm2(x))
        return x, attention

    def forward(self, x, y=None, mask=None):
        x = x + self.attn(self.norm1(x), y, mask)[0]
        x = x + self.mlp(self.norm2(x))
        return x

    def __init__(self, dim_self, dim_ref, num_heads, mlp_ratio=4., bias=False, dropout=0., act=nnf.relu,
                 norm_layer: nn.Module = nn.LayerNorm):
        super().__init__()
        self.norm1 = norm_layer(dim_self)
        self.attn = MultiHeadAttention(dim_self, dim_ref, num_heads, bias=bias, dropout=dropout)
        self.norm2 = norm_layer(dim_self)
        self.mlp = MlpTransformer(dim_self, int(dim_self * mlp_ratio), act=act, dropout=dropout)


class Transformer(nn.Module):

    def forward_with_attention(self, x, y=None, mask=None):
        attentions = []
        for layer in self.layers:
            x, att = layer.forward_with_attention(x, y, mask)
            attentions.append(att)
        return x, attentions

    def forward(self, x, y=None, mask=None):
        for i, layer in enumerate(self.layers):
            if i % 2 == 0 and self.enc_dec: # cross
                x = layer(x, y)
            elif self.enc_dec:  # self
                x = layer(x, x, mask)
            else:  # self or cross
                x = layer(x, y, mask)
        return x

    def __init__(self, dim_self: int, num_heads: int, num_layers: int, dim_ref: Optional[int] = None,
                 mlp_ratio: float = 2., act=nnf.relu, norm_layer: nn.Module = nn.LayerNorm, enc_dec: bool = False):
        super(Transformer, self).__init__()
        dim_ref = dim_ref if dim_ref is not None else dim_self
        self.enc_dec = enc_dec
        if enc_dec:
            num_layers = num_layers * 2
        layers = []
        for i in range(num_layers):
            if i % 2 == 0 and enc_dec:  # cross
                layers.append(TransformerLayer(dim_self, dim_ref, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))
            elif enc_dec:  # self
                layers.append(TransformerLayer(dim_self, dim_self, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))
            else:  # self or cross
                layers.append(TransformerLayer(dim_self, dim_ref, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))
        self.layers = nn.ModuleList(layers)


class TransformerMapper(nn.Module):

    def forward(self, x):
        x = self.linear(x).view(x.shape[0], self.clip_length, -1)
        prefix = self.prefix_const.unsqueeze(0).expand(x.shape[0], *self.prefix_const.shape)
        prefix = torch.cat((x, prefix), dim=1)
        """拼接的目的：
            将输入x(如CLIP图像特征)与可学习的prefix拼接, 使Transformer在处理序列时,
            能够通过自注意力机制动态捕获x与prefix的交互关系。这种交互允许prefix根据输入内容自适应调整,
            而非静态参数。"""
        out = self.transformer(prefix)[:, self.clip_length:]
        return out

    def __init__(self, dim_clip: int, dim_embedding: int, prefix_length: int, clip_length: int, num_layers: int = 8):
        super(TransformerMapper, self).__init__()
        self.clip_length = clip_length
        self.transformer = Transformer(dim_embedding, 8, num_layers)
        self.linear = nn.Linear(dim_clip, clip_length * dim_embedding)
        self.prefix_const = nn.Parameter(torch.randn(prefix_length, dim_embedding), requires_grad=True)


class ClipCaptionModel(nn.Module):

    def get_dummy_token(self, batch_size: int, device: torch.device) -> torch.Tensor:
        return torch.zeros(batch_size, self.prefix_length, dtype=torch.int64, device=device)

    def forward(self, tokens: torch.Tensor, prefix: torch.Tensor, mask: Optional[torch.Tensor] = None,
                labels: Optional[torch.Tensor] = None):
        embedding_text = self.gpt.transformer.wte(tokens)
        prefix_projections = self.clip_project(prefix).view(-1, self.prefix_length, self.gpt_embedding_size)
        embedding_cat = torch.cat((prefix_projections, embedding_text), dim=1)
        # [b, pre_len + token_len, embed_size]
        if labels is not None:
            dummy_token = self.get_dummy_token(tokens.shape[0], tokens.device)
            labels = torch.cat((dummy_token, tokens), dim=1)
        out = self.gpt(inputs_embeds=embedding_cat, labels=labels, attention_mask=mask)
        return out

    def __init__(self, prefix_length: int, clip_length: Optional[int] = None, prefix_size: int = 512,
                 num_layers: int = 8, mapping_type: MappingType = MappingType.MLP):
        super(ClipCaptionModel, self).__init__()
        self.prefix_length = prefix_length
        self.gpt = GPT2LMHeadModel.from_pretrained('gpt2')
        self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1]
        if mapping_type == MappingType.MLP:
            self.clip_project = MLP((prefix_size, (self.gpt_embedding_size * prefix_length) // 2,
                                     self.gpt_embedding_size * prefix_length))
        else:
            self.clip_project = TransformerMapper(prefix_size, self.gpt_embedding_size, prefix_length,
                                                                     clip_length, num_layers)


class ClipCaptionPrefix(ClipCaptionModel):

    def parameters(self, recurse: bool = True):
        return self.clip_project.parameters()

    def train(self, mode: bool = True):
        super(ClipCaptionPrefix, self).train(mode)
        self.gpt.eval()
        return self



class FFN(nn.Module):
    def __init__(self, dim_input, dim_hidden, dim_out, dropout=0.):
        super(FFN, self).__init__()
        self.linear1 = nn.Linear(dim_input, dim_hidden)
        self.act = nn.ReLU()
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_hidden, dim_out)

    def forward(self, x):
        x = self.dropout(self.act(self.linear1(x)))
        x = self.dropout(self.linear2(x))
        return x

class my_MultiHeadAttention(nn.Module):
    def __init__(self, dim_sel, dim_ref, num_head, bias=True, dropout=0.):
        super(my_MultiHeadAttention, self).__init__()
        self.head = num_head
        self.dim = dim_sel // num_head
        self.scale = self.dim ** -0.5
        self.linear1 = nn.Linear(dim_sel, dim_sel, bias=bias)
        self.linear2 = nn.Linear(dim_ref, dim_ref * 2, bias=bias)
        self.linear3 = nn.Linear(dim_sel, dim_sel)
        self.dropout = nn.Dropout(dropout)

    
    def forward(self, x, y = None, mask = None):
        y = x if y is None else y
        b, n, d_x = x.shape
        _, m, d_y = y.shape
        # x: [b, n, dim_embed]
        q = self.linear1(x).view(b, n, self.head, d_x // self.head).transpose(2, 1)
        # q: [b, h, n, d]
        kv = self.linear2(y).view(b, m, 2, self.head, d_y // self.head)
        k, v = kv[:, :, 0].transpose(2, 1), kv[:, :, 1].transpose(2, 1)
        # k, v: [b, h, m, d] 
        attention = torch.einsum("bhnd,bhmd->bhnm", q, k) * self.scale

        if mask is not None:
            # mask: [b, n]
            assert mask.shape[-1] == attention.shape[-1], "shape is not according"
            mask = mask[:, None, :, None]
            attention.masked_fill(~mask, float('-inf'))
        
        attention = attention.softmax(dim=-1)
        out = torch.einsum("bhnm,bhmd->bhnd", attention, v).view(b, n, d_x)
        out = self.dropout(self.linear3(out))
        return out, attention

class my_TransformerLayer(nn.Module):
    def __init__(self, dim_sel, dim_ref, num_head, mlp_rio=2.0):
        super(my_TransformerLayer, self).__init__()
        self.normlize1 = nn.LayerNorm(dim_sel)
        self.attn = my_MultiHeadAttention(dim_sel, dim_ref, num_head)
        self.normlize2 = nn.LayerNorm(dim_sel)
        self.ffn = FFN(dim_sel, int(dim_sel * mlp_rio), dim_sel)

    def forward(self, x, y=None, mask = None):
        x = x + self.attn(self.normlize1(x), y, mask)[0]
        x = x + self.ffn(self.normlize2(x))
        return x

class my_Transformer(nn.Module):
    def __init__(self, dim_sel, num_head, mlp_rio=2.0, dim_ref: Optional[int] = None, num_layer=8, enc_dec=False):
        super(my_Transformer, self).__init__()
        dim_ref = dim_ref if dim_ref is not None else dim_sel
        if enc_dec:
            num_layer = num_layer * 2
        layers = []
        for i in range(num_layer):
            if i % 2 == 0 and enc_dec:  # cross
                layer = my_TransformerLayer(dim_sel, dim_ref, num_head, mlp_rio)
            elif enc_dec:
                layer = my_TransformerLayer(dim_sel, dim_sel, num_head, mlp_rio)
            else:
                layer = my_TransformerLayer(dim_sel, dim_ref, num_head, mlp_rio)
            layers.append(layer)
        self.transformer = nn.ModuleList(layers)


class my_TransformerMapper(nn.Module):
    def __init__(self, prefix_dim, dim_embed: int, prefix_length, clip_length, num_layers=8):
        super(my_TransformerMapper, self).__init__()
        self.clip_length = clip_length
        self.linear = nn.Linear(prefix_dim, clip_length * dim_embed)
        self.project = my_Transformer(prefix_dim, 8, num_layer=num_layers)
        self.prefix_const = nn.Parameter(torch.randn((prefix_length, dim_embed)), requires_grad=True)

    def forward(self, x):
        x = self.linear(x).view(x.shape[0], self.clip_length, -1)
        prefix_extend = self.prefix_const.unsqueeze(0).expand(x.shape[0], *self.prefix_const.shape)
        prefix = torch.cat([x, prefix_extend], dim=1) # [batch_size, prefix_length*2, dim_embed]
        out = self.project(prefix)[:, self.clip_length:]
        return out

class my_MLP(nn.Module):
    def __init__(self, sizes: Tuple[int, ...], bias=True, act = nn.Tanh):
        layers = []
        for i in range(len(sizes) - 1):
            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias))
            if i < len(sizes) - 2:
                layers.append(act())
        self.model = nn.ModuleList(layers)
    def forward(self, x):
        return self.model(x)


class my_ClipCaptionModel(nn.Module):
    def __init__(self, clip_length, prefix_length, map_type = 'mlp', prefix_dim = 512, num_layer = 8):
        super(my_ClipCaptionModel, self).__init__()
        self.prefix_len = prefix_length
        self.gpt = GPT2LMHeadModel.from_pretrained('gpt2')
        self.embed_dim = self.gpt.transformer.wte.weight.shape[1]
        if map_type == 'mlp':
            self.project = my_MLP((prefix_dim, (prefix_length * self.embed_dim) // 2, prefix_length * self.embed_dim))
        else:
            self.project = my_TransformerMapper(prefix_dim, self.embed_dim, prefix_length, clip_length, num_layer)
    
    def forward(self, tokens, mask, prefix, labels = None):
        token_embeds = self.gpt.transformer.wte(tokens)
        prefix = self.project(prefix).view(-1, self.prefix_len, self.embed_dim)
        embed_cat = torch.cat([prefix, token_embeds], dim=1)
        out = self.gpt(inputs_embeds=embed_cat, labels=labels, attention_mask = mask)
        return out